import warnings

from .common.train import train
from .common.model import model
from .common.optimizer import optimizer
from .common.scheduler import lr_multiplier
from .common.dataloader import dataloader

bs = 16
n_gpus = 2
BS = bs * n_gpus

if BS != 32:
    warnings.warn(f'VitMatte Small version: BS={BS} is not 20')

train.max_iter = int(43100 / BS * 100)
train.checkpointer.period = int(43100 / BS * 10)

optimizer.lr = 5e-4
lr_multiplier.scheduler.values = [1.0, 0.1, 0.05]
lr_multiplier.scheduler.milestones = [int(43100 / BS * 30), int(43100 / BS * 90)]
lr_multiplier.scheduler.num_updates = train.max_iter
lr_multiplier.warmup_length = 250 / train.max_iter

train.init_checkpoint = './pretrained/dino_vit_s_fna.pth'
train.output_dir = './output_of_train/ViTMatte_S_100ep'

dataloader.train.batch_size = bs
dataloader.train.num_workers = n_gpus
